export NUM_NODES=1
export NUM_GPUS_PER_NODE=8


MODEL_SIZE="llama-7b"
MEMORY=34

MODEL_ARGS="
    --model_size ${MODEL_SIZE} \
    --set_model_config_manually 1 \
    --set_layernum_manually 1 \
    --vocab_size 32000 \
    --hidden_size 4096 \
    --num_hidden_layers 12 \
    --num_attention_heads 32 \
    --seq_length 256"

# lyh !!!!注意 如果要手动配置的话 要另set_model_config_manually = 1
# TODO 先测试0 0不自动的

BSZ_ARGS="
    --min_bsz 16 \
    --max_bsz 1024 \
    --bsz_scale 16 \
    --settle_bsz 64 \
    --settle_chunk 8 \
    --recommend_min_bsz 0
"

SEARCH_SPACE_ARGS="
    --search_space full \
    --disable_dp 0 \
    --disable_tp 0 \
    --disable_pp 0 \
    --disable_sdp 0 \
    --disable_ckpt 0 \
    --disable_vtp 0 \
    --disable_tp_consec 1 \
    --max_tp_deg 8 \
    --max_pp_deg 16 \
    --fine_grained_mode 0 \
    --computation_mode linear
"

# gees TODO settle_bsz 和 settle_chunk 和 mixed_precision改为FP32
# todo gees虽然可以接受global batch size用于以前全profile的自动并行，
#  但是在真实运行pretrain的时候接受的还是micro batch size，所以搜索的最后一部应当记录pp tp dp chunk=num of mb 以及micro batch size


SEARCH_ARGS="
    ${BSZ_ARGS} \
    ${SEARCH_SPACE_ARGS} \
    ${MODEL_ARGS} \
    --num_nodes ${NUM_NODES} \
    --num_gpus_per_node ${NUM_GPUS_PER_NODE} \
    --memory_constraint $MEMORY \
    --mixed_precision fp32 \
    --pipeline_type pipedream_flush \
    --default_dp_type zero2 \
    --embed_sdp 0
"

BACKGROUND=0

if [ $BACKGROUND -eq 1 ]; then
    echo "Search in background..."
    OUTPUT_FILE="Search_${MODEL_SIZE}_${MEMORY}GB_${NUM_NODES}Nodes_${NUM_GPUS_PER_NODE}GPUs_per_node.log"
    nohup python3 search_dist.py ${SEARCH_ARGS} 1> ${OUTPUT_FILE} 2>&1 &
else
    echo "Search in foreground..."
    python3 search_dist.py ${SEARCH_ARGS}
fi